package ru.scalabook.algorithms.heap

import scala.math.Ordered.orderingToOrdered

enum BinaryHeap[+A]:
  import BinaryHeap.*

  case Leaf

  private case Branch(
      min: A,
      left: BinaryHeap[A],
      right: BinaryHeap[A],
      override val size: Int,
      override val height: Int
  )

  lazy val isEmpty: Boolean = this match
    case Leaf => true
    case _    => false

  val size: Int = this match
    case neh: Branch[A] => neh.size
    case _              => 0

  val height: Int = this match
    case neh: Branch[A] => neh.height
    case _              => 0

  lazy val minOption: Option[A] = this match
    case neh: Branch[A] => Some(neh.min)
    case _              => None

  def insert[B >: A: Ordering](x: B): BinaryHeap[B] = this match
    case Leaf => BinaryHeap(x)
    case Branch(min, left, right, size, height) =>
      if left.size < left.height * left.height - 1 then
        bubbleUp(min, left.insert(x), right)
      else if right.size < right.height * right.height - 1 then
        bubbleUp(min, left, right.insert(x))
      else if right.height < left.height then
        bubbleUp(min, left, right.insert(x))
      else
        bubbleUp(min, left.insert(x), right)

  def remove[B >: A: Ordering]: Option[BinaryHeap[B]] =
    def floatLeft(x: A, l: BinaryHeap[A], r: BinaryHeap[A]): BinaryHeap[A] =
      l match
        case Branch(y, lt, rt, _, _) => BinaryHeap(y, BinaryHeap(x, lt, rt), r)
        case _                       => BinaryHeap(x, l, r)

    def floatRight(x: A, l: BinaryHeap[A], r: BinaryHeap[A]): BinaryHeap[A] =
      r match
        case Branch(y, lt, rt, _, _) => BinaryHeap(y, l, BinaryHeap(x, lt, rt))
        case _                       => BinaryHeap(x, l, r)

    def mergeChildren(l: BinaryHeap[A], r: BinaryHeap[A]): BinaryHeap[A] =
      (l, r) match
        case (Leaf, Leaf) => Leaf
        case (Leaf, Branch(rmin, rleft, rright, rsize, rheight)) =>
          floatRight(rmin, l, mergeChildren(rleft, rright))
        case (Branch(lmin, lleft, lright, lsize, lheight), Leaf) =>
          floatLeft(lmin, mergeChildren(lleft, lright), r)
        case (
              Branch(lmin, lleft, lright, lsize, lheight),
              Branch(rmin, rleft, rright, rsize, rheight)
            ) =>
          if lsize < lheight * lheight - 1 then
            floatLeft(lmin, mergeChildren(lleft, lright), r)
          else if rsize < rheight * rheight - 1 then
            floatRight(rmin, l, mergeChildren(rleft, rright))
          else if rheight < lheight then
            floatLeft(lmin, mergeChildren(lleft, lright), r)
          else floatRight(rmin, l, mergeChildren(rleft, rright))

    def bubbleRootDown(h: BinaryHeap[A]): BinaryHeap[A] = h match
      case Branch(min, left, right, _, _) =>
        bubbleDown(min, left, right)
      case _ => Leaf

    this match
      case Branch(_, left, right, _, _) =>
        Some(bubbleRootDown(mergeChildren(left, right)))
      case _ => None
  end remove
end BinaryHeap

object BinaryHeap:
  def empty[A]: BinaryHeap[A] = Leaf

  def apply[A: Ordering](x: A): BinaryHeap[A] =
    BinaryHeap(x, Leaf, Leaf)

  private def apply[A: Ordering](
      x: A,
      l: BinaryHeap[A],
      r: BinaryHeap[A]
  ): BinaryHeap[A] =
    Branch(x, l, r, l.size + r.size + 1, math.max(l.height, r.height) + 1)

  private def bubbleUp[A: Ordering](
      x: A,
      l: BinaryHeap[A],
      r: BinaryHeap[A]
  ): BinaryHeap[A] = (l, r) match
    case (Branch(y, lt, rt, _, _), _) if x > y =>
      BinaryHeap(y, BinaryHeap(x, lt, rt), r)
    case (_, Branch(z, lt, rt, _, _)) if x > z =>
      BinaryHeap(z, l, BinaryHeap(x, lt, rt))
    case (_, _) => BinaryHeap(x, l, r)

  private def bubbleDown[A: Ordering](
      x: A,
      l: BinaryHeap[A],
      r: BinaryHeap[A]
  ): BinaryHeap[A] = (l, r) match
    case (Branch(y, _, _, _, _), Branch(z, lt, rt, _, _))
        if z < y && z < x =>
      BinaryHeap(z, l, bubbleDown(x, lt, rt))
    case (Branch(y, lt, rt, _, _), _) if x > y =>
      BinaryHeap(y, bubbleDown(x, lt, rt), r)
    case (_, _) => BinaryHeap(x, l, r)
end BinaryHeap
